import torch 
import torchvision
import torch.nn.functional as F
import torchvision.transforms as transforms
import torch.optim as optim
import matplotlib.pyplot as plt
import numpy as np
import numpy as np
from PIL import Image
from torchvision import transforms
from torch.utils.tensorboard import SummaryWriter
#预处理
myTransforms = transforms.Compose(
    [transforms.RandomResizedCrop((40,40)),#切割
     transforms.RandomHorizontalFlip(),#翻转
     transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])#归一化
train_dataset = torchvision.datasets.ImageFolder(root='/home/lyx/Desktop/LYX_python/dataset/tldataset/train',transform=myTransforms)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=2)

test_dataset = torchvision.datasets.ImageFolder(root='/home/lyx/Desktop/LYX_python/dataset/tldataset/train', transform=myTransforms)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=4, shuffle=True, num_workers=2)

class myNetwork(torch.nn.Module):
    def __init__(self):
        super(myNetwork, self).__init__()
        self.conv1 = torch.nn.Conv2d(in_channels=3, out_channels = 32, kernel_size=5, stride=1, padding=2)
        self.pool1 = torch.nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = torch.nn.Conv2d(in_channels=32, out_channels = 32, kernel_size=5, stride=1, padding=2)
        self.pool2 = torch.nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv3 = torch.nn.Conv2d(in_channels=32, out_channels = 32, kernel_size=5, stride=1, padding=2)
        self.pool3 = torch.nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = torch.nn.Linear(in_features = 32*5*5, out_features = 64, bias=True)
        self.fc2 = torch.nn.Linear(in_features = 64, out_features = 64, bias=True)
        self.fc3 = torch.nn.Linear(in_features = 64, out_features = 10, bias=True)

    def forward(self, train_data):
        output = self.pool1(torch.nn.functional.relu(self.conv1(train_data)))
        output = self.pool2(torch.nn.functional.relu(self.conv2(output)))
        output = self.pool3(torch.nn.functional.relu(self.conv3(output)))
        output = output.view(-1, 32*5*5)#可视化输出
        output = torch.nn.functional.relu(self.fc1(output))
        output = torch.nn.functional.relu(self.fc2(output))
        output = self.fc3(output)

        return output

if __name__=="__main__":
    myModel = myNetwork()
    myOptimzier = optim.SGD(myModel.parameters(), lr = 0.01, momentum=0.9)
    myLoss = torch.nn.CrossEntropyLoss()

    for _epoch in range(10):
        training_loss = 0.0
        for _step, input_data in enumerate(train_loader):

            image, label = input_data
            predict_label = myModel.forward(image)
            loss = myLoss(predict_label, label)

        # myWriter.add_scalar('training loss', loss, global_step = _epoch*len(train_loader) + _step)

            myOptimzier.zero_grad()
            loss.backward()
            myOptimzier.step()

            training_loss = training_loss + loss.item()
            if _step % 10 == 0 :
                print('[iteration - %3d] training loss: %.3f' % (_epoch*len(train_loader) + _step, training_loss/10))
                training_loss = 0.0


    correct = 0
    total = 0
    with torch.no_grad():
        for images,labels in test_loader:
            outputs = myModel(images)
            numbers,predicted = torch.max(outputs.data,1)
            total += labels.size(0)
            correct += (predicted==labels).sum().item()

    print('Testing Accuracy : %d %%' % ( 100 * correct / total))




